# lama/optimization.py

import numpy as np
from scipy.optimize import minimize_scalar


def update_z(G, C, N, z, U, alpha=0.5, gamma=0.1):
    """
    更新每个网络中节点的归属值 z(w)(u)。
    :param G: 多层融合图
    :param C: 核心社区
    :param N: 边界节点集合
    :param z: 当前节点归属字典
    :param U: 统一从属关系
    :param alpha: 层内权重
    :param gamma: 正则化参数
    :return: 更新后的 z 字典
    """
    new_z = z.copy()
    for node in N:
        score = 0
        for neighbor in G[node]:
            if neighbor in C:
                score += G[node][neighbor]['weight'] * z.get(neighbor, 0.5)
        score /= (len(N) + 1e-9)
        new_z[node] = score
    return new_z


def update_U(z_list, U, delta_w=None):
    """
    统一不同网络的从属关系。
    :param z_list: 各网络的 z 归属字典列表
    :param U: 当前统一从属关系
    :param delta_w: 每个网络的权重
    :return: 新的统一从属关系
    """
    if delta_w is None:
        delta_w = [1.0 / len(z_list)] * len(z_list)

    new_U = {}
    all_nodes = set().union(*[z.keys() for z in z_list])

    for node in all_nodes:
        weighted_sum = sum(delta_w[i] * z_dict.get(node, 0.5) for i, z_dict in enumerate(z_list))
        new_U[node] = weighted_sum
    return new_U



def compute_similarity(z1, z2):
    """
    计算两个网络之间的节点归属相似性（用于层间一致性）
    """
    common_nodes = set(z1.keys()) & set(z2.keys())
    similarity = 0
    for node in common_nodes:
        similarity += abs(z1[node] - z2[node])
    return similarity / (len(common_nodes) + 1e-9)


def compute_combined_weight(G, u, v, delta_w):
    """
    计算综合边权重：
    w_combined(u,v) = α * w_ppi(u,v) + (1-α) * w_ge(u,v)
    """
    ppi_weight = G[u][v].get('PPI', {}).get('weight', 0)
    expr_weight = G[u][v].get('expr', {}).get('weight', 0)
    return delta_w['PPI'] * ppi_weight + delta_w['Expression'] * expr_weight


def objective_function(z, U, delta_w, G):
    """
    LAMA 的目标函数：
    Eq = ∑δw * ||zw - U||^2 + λ∑zwn^T L zw
    """
    loss = 0
    for node in U:
        loss += sum(delta_w[i] * (z[i].get(node, 0.5) - U[node]) ** 2 for i in range(len(z)))
    # 可选：加入 Laplacian 平滑项
    return loss


def learn_delta_weights(similarities):
    """
    自动学习每层的 δw 权重，基于相似性得分。
    """
    total_sim = sum(similarities.values())
    return {k: v / total_sim for k, v in similarities.items()}


def update_z_and_U(z_list, U, delta_w, lambda_reg=0.1):
    """
    更新每个网络的节点归属 z(w)(u) 和统一从属关系 U(u)
    :param z_list: 各网络层的 z 字典列表 [z_ppi, z_expr]
    :param U: 统一从属关系字典 {u: value}
    :param delta_w: 层权重 {'PPI': 0.6, 'Expression': 0.4}
    :param lambda_reg: 正则化参数，控制一致性
    :return: 新的 z 列表和新的 U 字典
    """
    # Step 1: 更新统一从属关系 U
    new_U = {}
    all_nodes = set()
    for z_dict in z_list:
        all_nodes.update(z_dict.keys())

    for node in all_nodes:
        weighted_sum = 0
        total_weight = 0
        for i, z_dict in enumerate(z_list):
            if node in z_dict:
                weighted_sum += delta_w[i] * z_dict[node]
                total_weight += delta_w[i]
        new_U[node] = weighted_sum / (total_weight + 1e-9) if total_weight else 0.5

    # Step 2: 使用新的 U 更新每个网络的 z(w)(u)
    new_z_list = []
    for layer_idx, z_dict in enumerate(z_list):
        new_z = {}
        for node, old_z in z_dict.items():
            # 计算目标函数梯度方向上的更新
            u_val = new_U.get(node, 0.5)
            grad_term = 2 * delta_w[layer_idx] * (u_val - old_z)
            new_z[node] = old_z + 0.1 * grad_term  # 学习率设为 0.1
        new_z_list.append(new_z)

    return new_z_list, new_U